import torch
#任务把字符串hello输出为ohllo
input_size = 4
num_class = 4
hidden_size = 8
#嵌入到10层
embedding_size = 10
batch_size = 1
#两层rnn
num_layers = 2
seq_len = 5
#创建词表
idx2char_1 = ['e', 'h', 'l', 'o']
idx2char_2 = ['h', 'l', 'o']

x_data = [[1, 0, 2, 2, 3]]
y_data = [3, 1, 2, 2, 3]

# inputs 作为交叉熵中的Inputs，维度为（batchsize，seqLen）(1,5)
inputs = torch.LongTensor(x_data)
# labels 作为交叉熵中的Target，维度为（batchsize*seqLen）(1,5)
labels = torch.LongTensor(y_data)

#===================================================embedding学习========================================================
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        #emb把独热向量变成embedding的向量优化 其维度为(input_size, embedding_size) (4,10) 输入的每个向量维度为4 一共输入5个向量
        self.emb = torch.nn.Embedding(input_size, embedding_size)
        """
        input_size我们的输入维度原来是一个四维向量 但经过embedding前向传播变成十维了 所以input_size=embedding_size即十维度
        hidden_size隐层输出大小 即8
        num_layers即rnn层数
        """
        self.rnn = torch.nn.RNN(input_size=embedding_size,
                                hidden_size=hidden_size,
                                num_layers=num_layers,
                                batch_first=True)

        self.fc = torch.nn.Linear(hidden_size, num_class)

    #x为(batch,seq) (1,5)
    def forward(self, x):
        #第0层隐层输入为全0 (2,1,8)
        hidden = torch.zeros(num_layers, x.size(0), hidden_size)
        #embedding前向传播 (batch,seq,embedding) (1,5,10)
        """
        1.embedding层会对输入的数据进行查看 然后选择对应的行
        2.例如当embedding为[4,10] x为[1,5]输入为[[0,1,2,0,1]] 则执行embedding(x)时会依次选择embedding矩阵的第0,1,2,0,1行
        组成一个新矩阵即[1,5,10]
        3.我们的x不能超出embedding的行数 比如以上面的例子为例x为[1,5]输入为[[0,1,2,0,9]] 9超出了[0,3]会报错
        """
        x = self.emb(x)
        #x通过hidden隐层变成[1,5,8]
        x, _ = self.rnn(x, hidden)
        #x通过全连接层从[1,5,8]-->[1,5,4]
        x = self.fc(x)
        return x.view(-1, num_class)


net = Model()

criterion = torch.nn.CrossEntropyLoss()
#反向传播也会跟新embedding参数
optimizer = torch.optim.Adam(net.parameters(), lr=0.05)
#net.state_dict() 模型的参数(不包括结构)以字典形式返回
print("00训练前",net.state_dict())
for epoch in range(15):
    optimizer.zero_grad()
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    _, idx = outputs.max(dim=1)
    idx = idx.data.numpy()
    print('Predicted string: ', ''.join([idx2char_1[x] for x in idx]), end='')
    print(", Epoch [%d/15] loss = %.3f" % (epoch + 1, loss.item()))

#==================================================模型保存==============================================================
netname  = "myNet"
#后缀是什么都可以 推荐.pth
file1 = f"checkpoint/{netname}-{embedding_size}.pth"
file2 = f"checkpoint/{netname}-{embedding_size}-dict.pth"
#保存又叫序列化把内存中的数据变成0和1保存在硬盘中 加载又称反序列化把硬盘中的0和1数据重新加载回内存变成对象
#保存方式1 保存到指定文件 save保存的是模型结构+模型参数 文件较大 torch.save遇到同名文件直接覆盖
torch.save(net,file1)
#保存方式2 state_dict()保存的是模型参数是以字典的形式保存的 只保留参数文件所以比保存方式1的文件小
torch.save(net.state_dict(),file2)
#读取方法1 和保存方法1对应 直接使用load函数
#读取方法1对应保存方法1 因为保存了结构所以不需要实例化一个空的网络结构
model1 = torch.load(file1)
#读取方法2 和保存方法2对应 第一步使用load函数 第二部实例化你保存的模型 第三调用实例化后的模型load_state_dict方法
data2  = torch.load(file2)
#因为没有保存结构所以需要实例化一个空的网络结构
model2 = Model()
model2.load_state_dict(data2)
#推荐使用读取方法2和保存方法2
print("01保存方法1model1",model1)
print("02保存方法2model2",model2)


#==============================================模型保存进阶(训练中定期保存)=================================================
#训练模型时有可能会意外终止 所以在训练时有必要随时保存 这里默认每50个epoch保存一次
for epoch in range(120):
    optimizer.zero_grad()
    outputs = model2(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    #每过50个epoch保存模型参数 优化器参数 epoch参数 loss损失其中模型和优化器是重点因为随时会变化
    if(epoch + 1) % 50 ==0 :
        checkpoint = {
            "model_state+dict": model2.state_dict(),
            "optimizer_state_dict":optimizer.state_dict(),
            "epoch":epoch,
            "loss":loss.item()
        }
    #模拟意外中断
    if(epoch>105):
        print("中断了")
        break
        file3=f"checkpoint/{netname}-{embedding_size}-{epoch}.pth"
        torch.save(checkpoint,file3)
    _, idx = outputs.max(dim=1)
    idx = idx.data.numpy()
    print('Predicted string: ', ''.join([idx2char_1[x] for x in idx]), end='')
    print(", Epoch [%d/120] loss = %.3f" % (epoch + 1, loss.item()))

#中断恢复
file4 = "checkpoint/myNet-10-99.pth"
backup = torch.load(file4)

model3 = Model()
model3.load_state_dict(backup['model_state+dict'])

optimizer3 = torch.optim.Adam(net.parameters(), lr=0.05)
optimizer3.load_state_dict(backup['optimizer_state_dict'])

start_epoch = backup['epoch']
for epoch in range(start_epoch,120):
    optimizer.zero_grad()
    outputs = model3(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer3.step()

    _, idx = outputs.max(dim=1)
    idx = idx.data.numpy()
    print('Predicted string: ', ''.join([idx2char_1[x] for x in idx]), end='')
    print(", Epoch [%d/120] loss = %.3f" % (epoch + 1, loss.item()))

#=====================================================利用GPU训练========================================================
if(torch.cuda.is_available()):
    #模型放入cuda
    model3 = model3.cuda()
    #损失函数放入cuda
    loss = loss.cuda()
    #优化器没有cuda方法放入不了




